from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
from sklearn import tree
from sklearn import metrics
import os

# 使用IRIS数据集进行决策树例子

iris = load_iris()
# print(iris)
# print(iris['data'])

# 数据预处理
train_data, test_data, train_traget, test_target = train_test_split(
    iris['data'],
    iris['target'],
    test_size=0.2,  # 测试数据占比
    random_state=1  # 随机分配
)
# print(train_data)
# print(test_data)
# print(train_traget)
# print(test_target)


# 建立分类策略树模型
clf = tree.DecisionTreeClassifier(
    criterion="entropy"
)

# 使用训练数据进行训练
clf.fit(train_data, train_traget)

# 进行预测
y_pred = clf.predict(test_data)

# 进行验证
score = metrics.accuracy_score(
    y_true=test_target,
    y_pred=y_pred
)  # 获取准确度
print(score)

cfmt = metrics.confusion_matrix(
    y_true=test_target,
    y_pred=y_pred
)  # 获取结果混淆矩阵
print(cfmt)

# 可以对决策树进行输出分析
if not os.path.exists("./data"):
    os.makedirs("./data")
with open("./data/tree.txt", "w") as fw:
    tree.export_graphviz(clf, out_file=fw)
